#include "./googletest-1.17.0/googletest/include/gtest/gtest.h"
#include "./googletest-1.17.0/googlemock/include/gmock/gmock.h"

#include "../../loadBalance/code/RoundRobin.hpp"
#include "../../loadBalance/code/RandomSelect.hpp"
#include "../../loadBalance/code/WeightedRandom.hpp"
#include "../../loadBalance/code/HashSelect.hpp"
#include "../../loadBalance/code/ConsistentHash.hpp"

// ---------- RoundRobin ----------
TEST(RoundRobinTest, Sequential)
{
  std::vector<std::string> nodes{"Server1:8080", "Server2:8081", "Server3:8082"};
  LoadBalance::RoundRobin<std::string> rr(nodes);

  EXPECT_EQ(rr.getNext(), "Server1:8080");
  EXPECT_EQ(rr.getNext(), "Server2:8081");
  EXPECT_EQ(rr.getNext(), "Server3:8082");
  EXPECT_EQ(rr.getNext(), "Server1:8080");
  EXPECT_EQ(rr.getNext(), "Server2:8081");
}

// ---------- RandomSelect ----------
TEST(RandomSelectTest, AllResultsFromSet)
{
  std::vector<std::string> nodes{"Server1:8080", "Server2:8081", "Server3:8082"};
  LoadBalance::RandomSelect<std::string> rs(nodes);

  for (int i = 0; i < 10; ++i)
  {
    EXPECT_THAT(nodes, ::testing::Contains(rs.getNext())); // nodes中有Contais()
  }
}

// ---------- WeightedRandom ----------
TEST(WeightedRandomTest, AllResultsFromSet)
{
  std::vector<std::pair<std::string, int>> nodes = {{"Server1:8080", 1}, {"Server2:8081", 2}, {"Server3:8082", 3}};
  LoadBalance::WeightedRandom<std::string> wr(nodes);

  std::vector<std::string> names{}; // 提取名称比较
  for (const auto &n : nodes)
  {
    names.push_back(n.first);
  }
  for (int i = 0; i < 10; ++i)
  {
    EXPECT_THAT(names, ::testing::Contains(wr.getNext()));
  }

  // 权重固定，但随机选，不能测统计分布
  // // 取 6000 次，统计分布
  // std::map<std::string, int> count;
  // for (int i = 0; i < 6000; ++i)
  // {
  //   ++count[wr.getNext()];
  // }

  // // 权重 1:2:3，期望比例约 1/6 : 1/3 : 1/2
  // EXPECT_NEAR(count["Server1:8080"], 1000, 150); // 1/6 * 6000 ± 150
  // EXPECT_NEAR(count["Server2:8081"], 2000, 200); // 1/3 * 6000 ± 200
  // EXPECT_NEAR(count["Server3:8082"], 3000, 250); // 1/2 * 6000 ± 250
}

// ---------- HashSelect ----------
TEST(HashSelectTest, Deterministic)
{
  std::vector<std::string> nodes{"Server1:8080", "Server2:8081", "Server3:8082"};
  LoadBalance::HashSelect<std::string> hs(nodes);

  EXPECT_EQ(hs.getNode("key1"), hs.getNode("key1")); // 相同key落到同一节点
  EXPECT_EQ(hs.getNode("key2"), hs.getNode("key2"));
  EXPECT_EQ(hs.getNode("key3"), hs.getNode("key3"));
}

// ---------- ConsistentHash ----------
class ConsistentHashTest : public ::testing::Test // 多份TEST_F()相同初始化
{
protected:
  void SetUp() override
  {
    this->nodes_ = {"Server1:8080", "Server2:8081", "Server3:8082"};
    for (const auto &n : this->nodes_)
    {
      this->ch_.addNode(n);
    }
  }

  std::vector<std::string> nodes_;
  LoadBalance::ConsistentHash<std::string> ch_{3};
};

TEST_F(ConsistentHashTest, GetNode)
{
  std::vector<std::string> keys = {"key1", "key2", "key3", "key4", "key5"};
  for (const auto &k : keys)
  {
    EXPECT_THAT(this->nodes_, ::testing::Contains(this->ch_.getNode(k)));
  }
}

TEST_F(ConsistentHashTest, RemoveNode)
{
  std::vector<std::string> keys = {"key1", "key2", "key3", "key4", "key5"};

  // 删除前的映射
  std::map<std::string, std::string> before{};
  for (const auto &k : keys)
  {
    before[k] = this->ch_.getNode(k);
  }

  // 删除
  this->ch_.delNode("Server2:8081");

  // 删除后的映射，应只有Server1和Server3
  std::set<std::string> after{"Server1:8080", "Server3:8082"};
  for (const auto &k : keys)
  {
    EXPECT_THAT(after, ::testing::Contains(this->ch_.getNode(k)));
  }

  // 验证单调性：删除Server2后
  for (const auto &k : keys)
  {
    if (before[k] != "Server2:8081")
    {
      // before[k]：删除前某个k映射Server1或Server3 == 删除后某个k映射Server1或Server3，映射不变
      EXPECT_EQ(before[k], this->ch_.getNode(k)) << "key=" << k; // 若出错，打印key值
    }
  }
}

int main(int argc, char **argv)
{
  ::testing::InitGoogleTest(&argc, argv);
  return RUN_ALL_TESTS();
}